# coding: utf8

import sys
import json
import tensorflow as tf
from cnn import CNN
from tensorflow.examples.tutorials.mnist import input_data


class CG:

    def __init__(self, model_string, file_path):
        self.sess = tf.InteractiveSession()
        self.model = {}
        self._output = None

        self.model_json = json.loads(model_string)
        self.file_path = file_path

        for node in self.model_json["nodeDataArray"]:
            self.add_node(node)

        # for link in self.model_json["linkDataArray"]:
        #     print(link["from"], link["to"])

    def add_node(self, node):
        n_type = node["text"]
        n_key = node["key"]

        if n_type == "input":
            self.model["x"] = tf.placeholder(tf.float32, [None, 784])
            self.model["y_"] = tf.placeholder(tf.float32, [None, 10])
            self.model[n_key] = tf.reshape(self.model["x"], [-1, 28, 28, 1])
        elif n_type == "conv_2d":
            self.model[n_key] = CNN.conv2d(self._output, **node)
        elif n_type == "max_pool_2x2":
            self.model[n_key] = CNN.max_pool_2x2(self._output)
        elif n_type == "flat":
            self.model[n_key] = CNN.flat(self._output, node["xdim"])
        elif n_type == "full_connection":
            self.model[n_key] = CNN.full_connection(self._output, **node)
        elif n_type == "softmax":
            self.model[n_key] = CNN.softmax(self._output)
        elif n_type == "dropout":
            self.model["keep_prob"] = tf.placeholder(tf.float32)
            self.model[n_key] = CNN.dropout(self._output, self.model["keep_prob"])

        self._output = self.model[n_key]

    def exec(self):
        mnist = input_data.read_data_sets(self.file_path, one_hot=True)

        # cost
        cross_entropy = tf.reduce_mean(-tf.reduce_sum(self.model["y_"] * tf.log(self._output),
                                                      reduction_indices=[1]))
        train_step = tf.train.AdamOptimizer(1e-4).minimize(cross_entropy)

        # accuracy
        correct_prediction = tf.equal(tf.argmax(self._output, 1), tf.argmax(self.model["y_"], 1))
        accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

        tf.global_variables_initializer().run()
        for i in range(1000):
            batch = mnist.train.next_batch(50)
            if i % 100 == 0:
                train_accuracy = accuracy.eval(feed_dict={
                    self.model["x"]: batch[0],
                    self.model["y_"]: batch[1],
                    self.model["keep_prob"]: 1.0})
                print(f"step {i}, training accuracy {train_accuracy}", flush=True)

            train_step.run(feed_dict={
                self.model["x"]: batch[0],
                self.model["y_"]: batch[1],
                self.model["keep_prob"]: 0.5})

        print(f"Final test accuracy %g" % accuracy.eval(feed_dict={
            self.model["x"]: mnist.test.images,
            self.model["y_"]: mnist.test.labels,
            self.model["keep_prob"]: 1.0
        }), flush=True)


if __name__ == "__main__":
    model_string = sys.argv[1]
    file_path = sys.argv[2]

    # model_string = open("/Users/jochen/Documents/lab/myth/visualization/cnn_mnist.json", encoding="utf8").read()
    # parameters = "/Users/jochen/Documents/lab/myth/Myth/Myth/data/mnist_data/"
    cg = CG(model_string, file_path)
    cg.exec()
